import logging
import os
import sqlite3
import json
import time
import jieba
import jieba.posseg as pseg
from docx import Document
import PyPDF2
import threading
import psycopg2
import tempfile


class Searcher:

    def __init__(self):
        self.conn = None
        self.cur = None
        self.init()
        self.is_connected = False
        self.psql_conn = None
        self.psql_cur = None
        with open('./psql_config.json', 'r') as f:
            self.psql_config = json.load(f)
        self.get_monitor_person()
        # logging.warning(self.monitor_list)

        with tempfile.NamedTemporaryFile(mode='w', delete=False) as temp_file:
            temp_file.write(self.monitor_list)
            temp_file_path = temp_file.name

        jieba.load_userdict(temp_file_path)
        try:
            os.unlink(temp_file_path)
        except:
            pass

    def psql_connect(self):
        try:
            self.psql_conn = psycopg2.connect(
                dbname=self.psql_config.get('dbname'),
                user=self.psql_config.get('user'),
                password=self.psql_config.get('password'),
                host=self.psql_config.get('host'),
                port=self.psql_config.get('port')
            )
            self.psql_cur = self.psql_conn.cursor()
        except:
            pass

    def psql_close(self):
        try:
            self.psql_conn.close()
        except:
            pass

    def sqlite_connect(self):
        try:
            self.conn = sqlite3.connect('document_index.db')
            self.cur = self.conn.cursor()
        except:
            pass

    def sqlite_close(self):
        try:
            self.conn.close()
        except:
            pass

    def init(self):
        self.sqlite_connect()
        # 创建表格
        self.cur.execute('''CREATE TABLE IF NOT EXISTS documents (
                        id INTEGER PRIMARY KEY AUTOINCREMENT,
                        path TEXT,
                        contents VARCHAR(500)
                    );''')

        # 创建索引
        self.cur.execute('CREATE INDEX IF NOT EXISTS idx_contents ON documents(contents)')
        self.conn.commit()
        self.sqlite_close()

    def insert_or_update_document(self, file_path, token):
        # 检查数据库中是否已存在具有相同 token 的记录
        self.cur.execute("SELECT * FROM documents WHERE contents = ?", (token,))
        existing_doc = self.cur.fetchone()

        if existing_doc:
            # 如果存在相同 token 的记录，则更新 file_path
            existing_path = existing_doc[1]
            existing_path_list = existing_path.split(';')
            existing_path_list.append(file_path)
            new_path = ";".join(set(existing_path_list))
            self.cur.execute("UPDATE documents SET path = ? WHERE contents = ?", (new_path, token))
        else:
            # 如果不存在相同 token 的记录，则插入新记录
            self.cur.execute("INSERT INTO documents (path, contents) VALUES (?, ?)", (file_path, token))

        self.conn.commit()

    # sqlite 不支持多线程
    def index_documents_all_threads(self, folder_path):
        path = sorted([os.path.join(folder_path, name) for name in os.listdir(folder_path) if os.path.isdir(os.path.join(folder_path, name))])
        path_split = [path[i:i+4] for i in range(0, len(path), 4)]
        threads = []
        for i in range(len(path_split)):
            thread = threading.Thread(target=self.batches_index_documents_all, args=(path_split[i],))
            threads.append(thread)
        for thread in threads:
            thread.start()
        for thread in threads:
            thread.join()

    def batches_index_documents_all(self, path_list):
        for path in path_list:
            self.index_documents_all(path)

    def index_documents_all(self, folder_path):
        for root, dirs, files in os.walk(folder_path):
            for file in files:
                file_path = os.path.join(root, file)
                self.index_documents_one(file_path)

    def index_documents_one(self, file_path):
        file_extension = os.path.splitext(file_path)[1].lower()
        if file_extension not in ['.txt', '.docx', '.pdf']:
            return

        try:
            # 读取文本文件内容
            if file_extension == '.txt':
                with open(file_path, 'r') as f:
                    content = f.read()

            # 读取 docx 文件内容
            elif file_extension == '.docx':
                doc = Document(file_path)
                content = ''
                for paragraph in doc.paragraphs:
                    content += paragraph.text

            # 读取 PDF 文件内容
            elif file_extension == '.pdf':
                with open(file_path, 'rb') as f:
                    reader = PyPDF2.PdfReader(f)
                    content = ''
                    for page_num in range(len(reader.pages)):
                        page = reader.pages[page_num]
                        content += page.extract_text()

            if 'Origin: Ubuntu' not in content:
                words = pseg.cut(content)
                person_names = set([word for word, flag in words if flag == 'nr'])
                self.sqlite_connect()
                for names in person_names:
                    self.insert_or_update_document(file_path, names)
                self.sqlite_close()
        except Exception as e:
            print(f"Error processing {file_path}: {e}")

    # 检索文档路径
    def search_documents(self, keyword):
        self.sqlite_connect()
        self.cur.execute('SELECT path FROM documents WHERE contents LIKE ?', ('%' + keyword + '%',))
        results = self.cur.fetchall()
        self.sqlite_close()
        return [result[0] for result in results]

    def get_monitor_person(self):
        self.psql_connect()
        sql = """
        SELECT name FROM risk_monitoring_list WHERE is_enable = 1
        """
        self.psql_cur.execute(sql)
        persons = set(self.psql_cur.fetchall())
        self.psql_close()
        self.monitor_list = ""
        for person in persons:
            self.monitor_list += f"{person[0]} nr\n"

    def __del__(self):
        try:
            self.conn.close()
            self.psql_conn.close()
        except:
            pass


if __name__ == '__main__':
    css = Searcher()
    start_time = time.time()
    css.index_documents_all('./marked_files/')
    # css.index_documents_all_threads('./marked_files/')
    end_time = time.time()
    print(f'{end_time - start_time}')

    keyword = input('请输入要检索的关键词： ')
    results = css.search_documents(keyword)

    if results:
        print('检索到以下文档路径：')
        for result in results:
            print(result)
    else:
        print('未找到匹配的文档路径。')
